Add seq parallelism for attention and MoE MLP #1328
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
suexu1025
requested review from
gobbleturk,
khatwanimohit,
bvandermoon,
RissyRan,
richjames0,
rni418 and
gagika
as code owners
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
suexu1025
changed the title
suexu1025
requested review from
shralex,
yangyuwei,
SurbhiJainUSC,
hengtaoguo and
A9isha
as code owners
RissyRan
reviewed
Mar 7, 2025
There was a problem hiding this comment.
Thanks for the change and improvement! Overall, LGTM! One thing I'd like to see how is the performance impact for training. When it's ready, could you help run a benchmark on 8X7b (or other model size) with FSDP + EP sharding in dropping (with and without this change)? Capturing profiles will be great! Thank you!
RissyRan
reviewed
Mar 7, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
with sp+ep, moe customer 2k seq inference improved by 20%
FIXES: b/374773995
Tests
tested on v6e/v5p:
SEQ=2048
python MaxText/inference_microbenchmark.py MaxText/configs/inference.yml max_prefill_predict_length=$SEQ max_target_length=6144 model_name=mixtral-8x7b ici_fsdp_parallelism=1 ici_autoregressive_parallelism=1 ici_expert_parallelism=1 ici_context_parallelism=4 ici_tensor_parallelism=1 scan_layers=false per_device_batch_size=1 attention=dot_product megablox=False quantization=int8 checkpoint_is_quantized=True quantize_kvcache=True capacity_factor=1 tokenizer_path=assets/tokenizer.mistral-v3 compute_axis_order=0,2,1,3 ar_cache_axis_order=0,2,1,3 enable_jax_profiler=True inference_microbenchmark_prefill_lengths="$SEQ" base_output_directory=$OUT_DIR run_name=$RUN_NAME profiler=xplane model_call_mode=inference inference_microbenchmark_stages=prefill
Checklist
Before submitting this PR, please make sure (put X in square brackets):